##############################################################################
import pandas
import pandas as pd
import numpy as np
def generate_distribution_mapping(
        df : pandas.core.frame.DataFrame, 
        origin_score : str, 
        target_score : str,
        range_left_end: int,
        range_right_end: int,
        ) -> dict:
    """
    Use the percentile mapping from the original distribution to the target 
    distribution.
    
    This is the second attempt, which scales the target score series to the 
    same length as the origin score series and establishes the mapping between
    two score series. If a score from the origin score series has multiple
    matched scores from the target score series, this method assgin each 
    matched target score with a probability based on its frequency. Otherwise,
    the origin score will be assign a single target score.

    Parameters
    ----------
    df : pandas.core.frame.DataFrame
        The dataframe that contains the original score series and the target
        score series.
    origin_score : str
        The column name of the original score series in the given dataframe.
    target_score : str
        The column name of the original target series in the given dataframe.
    range_left_end : int
        The left end point of the range of the origin score series.
    range_right_end : int
        The right end point of the range of the origin score series.

    Returns
    -------
    distribution_mapping : dict
        The mapping dictionary from the origin distribution to the target 
        distribution.

    Examples
    -------
    Convert the distribution of score series to the distribution of the
    target score series.
    
    >>> distribution_mapping = generate_distribution_mapping(
    ...     df=data_all, origin_score='origin', target_score='target',
    ...     range_left_end=300, range_right_end=1000
    ...     )
    >>> distribution_mapping
    {977.0: 897.0,
      978.0: {'score': [899.0, 897.0],
      'prob': [0.6666666666666666, 0.3333333333333333]}}
    """
    
    # score 需要拟合的分数
    # target 拟合的目标分布
    data_total = df[[origin_score, target_score]]
    data_total.columns = ['score', 'target']
    data_total.loc[:,'target_2'] = data_total['target']
    
    # 转化为整数
    data_total.loc[:,'score'] = data_total['score'].apply(lambda x:x if np.isnan(x) else int(x))
    data_total.loc[:,'target'] = data_total['target'].apply(lambda x:x if np.isnan(x) else int(x))
    
    order = data_total[data_total['score'].notna()]['score'].sort_values()
    # 统计两个分布中，每个分数取值的个数， 升序排列 （相当于求分位数）
    origin_distribution = order.value_counts().sort_index()
    target_distribution = data_total[data_total['target'].notna()].groupby('target').count().sort_index()['target_2']
    
    # 两个分布的分数个数， 去除空值
    N_origin, N_target = origin_distribution.sum(), target_distribution.sum()
    
    # 将原分布按分数数量比例缩放至新分布的分数数目
    scaled_target_distribution = target_distribution / N_target * N_origin
    scaled_order = []
    remains = 0
    for target, num in scaled_target_distribution.items():
        score_num = remains + num
        scaled_order += [target] * int(score_num)
        remains = score_num - int(score_num)
    if len(scaled_order) != N_origin:
        scaled_order += [target_distribution.index[-1]]
    
    distribution_mapping = {} # 存储一个映射，从原分布的某分数，指向 目标分布的 对应分位数的分数
    for score_origin, score_detail in pd.DataFrame({'origin':order.values, 'target':scaled_order}).groupby('origin'):
        score_count = score_detail['target'].value_counts()
        if len(score_count) > 1:
            score_rate = score_count/score_detail.shape[0]
            distribution_mapping[score_origin] = {
                'score':score_rate.index.tolist(),
                'prob':score_rate.values.tolist()
                }
        else:
            distribution_mapping[score_origin] = score_count.index[0]
        
    # 将剩余未做映射的值分配相应映射
    last = data_total['target'].min()
    for i in range(range_left_end, range_right_end + 1):
        if i not in distribution_mapping:
            distribution_mapping[i] = last
        else:
            last = distribution_mapping[i]
            
    return distribution_mapping
        

#################################################################################################
def apply_mappling(
        origin_score : int or float, 
        mapping : dict, 
        seed : int = None
        ) -> int or float:
    """
    Use the generated mapping dictionary to generate the new score.

    Parameters
    ----------
    origin_score : int or float
        The original score that provided.
    mapping : dict
        The mapping from origin to target, could have rates.
    seed : int, optional
        Provide a random seed for reproducing the result. The default is None.

    Returns
    -------
    new_score : int or float
        The new score.

    Examples
    -------
    Convert the distribution of score series to the distribution of the
    target score series.
    
    >>> distribution_mapping
    {977.0: 897.0,
     978.0: {'score': [899.0, 897.0],
      'prob': [0.6666666666666666, 0.3333333333333333]}}
    
    >>> data_total['origin']
    0        977.0
    1        978.0
    2        978.0
    Name: score, Length: 3, dtype: float64
    
    >>> data_total[['id', 'origin']].apply(
    ...     lambda x:apply_mappling(x['origin'], distribution_mapping, int(hash(x['id']) % 1e9))
    ...         if not np.isnan(x['origin']) else x['origin'], axis=1)
    0        977.0
    1        978.0
    2        978.0
    Name: score, Length: 3, dtype: float64
    """
    new_score = mapping[origin_score]
    # if isinstance(new_score, float):
    #     return new_score
    if isinstance(new_score, dict):
        np.random.seed(seed)
        new_score = np.random.choice(a=new_score['score'], size=1, p=new_score['prob'])[0]
        return new_score
    else:
        return new_score